from read_CloudSat import reader
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
import pandas as pd
import netCDF4 as nc
from scipy.spatial import cKDTree

if __name__ == '__main__':
    # 读取CloudSat数据
    fname = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_GEOPROF/cloudsat_2020_1_4/2020007041321_72950_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
    f = reader(fname)
    lon, lat, elv = f.read_geo()
    data = f.read_sds('Radar_Reflectivity')
    height = f.read_sds('Height')
    time = f.read_time(datetime=True)  # 将时间转换为日期和时间格式07050000
    f.close()

    # # 从文件名提取开始和结束时间
    # filename = "/mnt/datastore/liudddata/fy_4Adata/FY202001_04/fy_cth_202001_04/FY4A-_AGRI--_N_DISK_1047E_L2-_CTH-_MULT_NOM_20200107050000_20200107051459_4000M_V0001.NC"
    # cth_file = nc.Dataset(filename, 'r')
    # cth_data = cth_file.variables['CTH'][:]
    # latitudes = []
    # longitudes = []
    # cth_values = []   # 从文件名提取开始和结束时间
    filename = "/mnt/datastore/liudddata/result/20200104new/2020010705_predicted_2d.nc"
    cth_file = nc.Dataset(filename, 'r')
    cth_data = cth_file.variables['cth'][:]
    cbh_data= cth_file.variables['predicted'][:]
    latitudes = []
    longitudes = []
    cth_values = []
    cbh_values = []


    # start_datetime, end_datetime = [datetime.strptime(x, "%Y%m%d%H%M%S") for x in filename.split('_')[14:16]]
    # 手动指定开始时间和结束时间
    start_datetime = datetime(2020, 1, 7, 5, 0, 0)
    end_datetime = datetime(2020, 1, 7, 5, 15, 0)
    # 寻找时间范围内的索引
    time_series = pd.to_datetime(time)
    start_index, end_index = np.searchsorted(time_series, [start_datetime, end_datetime])

    # 选择时间范围内的数据
    rad_refl_t = np.transpose(data[start_index:end_index])
    rad_refl_t[rad_refl_t <= -20] = np.nan  # 过滤数据

    # 获取对应时间范围内的高度数据，单位转换为km
    hgt = np.mean(height[start_index:end_index], axis=0) / 1000
    lat_t = lat[start_index:end_index]
    lon_t = lon[start_index:end_index]
    lat_t[np.isnan(lat_t)] = np.nan
    lon_t[np.isnan(lon_t)] = np.nan

    # 加载地理坐标数据
    coord_file_name = 'FY4A_coordinates.nc'
    coord_file_open = nc.Dataset(coord_file_name, 'r')
    lat_fy = coord_file_open.variables['lat'][:, :].T
    lon_fy = coord_file_open.variables['lon'][:, :].T
    lat_fy[np.isnan(lat_fy)] = -90.
    lon_fy[np.isnan(lon_fy)] = 360.
    coord_file_open.close()

    # 将 lat_fy 和 lon_fy 转换为笛卡尔坐标系以便使用 cKDTree
    fy_coords = np.deg2rad(np.vstack([lat_fy.ravel(), lon_fy.ravel()]).T)
    fy_coords_cartesian = np.array([np.cos(fy_coords[:, 0]) * np.cos(fy_coords[:, 1]),
                                    np.cos(fy_coords[:, 0]) * np.sin(fy_coords[:, 1]),
                                    np.sin(fy_coords[:, 0])]).T

    # 确保 fy_coords_cartesian 中没有 NaN 或无穷大/无穷小值
    is_finite = np.isfinite(fy_coords_cartesian).all(axis=1)

    # 仅保留所有值都是有限的行
    fy_coords_cartesian_clean = fy_coords_cartesian[is_finite]

    # 现在可以安全地创建 cKDTree
    tree = cKDTree(fy_coords_cartesian_clean)

    # 对于每个 lat_t 和 lon_t，找到2km范围内的 lat_fy 和 lon_fy
    radius = 2 / 6371  # 2km 转换为弧度
    for lat, lon in zip(lat_t, lon_t):
        point = np.deg2rad(np.array([lat, lon]))
        point_cartesian = np.array([np.cos(point[0]) * np.cos(point[1]),
                                    np.cos(point[0]) * np.sin(point[1]),
                                    np.sin(point[0])])

        # 执行半径搜索
        indices = tree.query_ball_point(point_cartesian, r=radius)

        # 处理找到的点
        if indices:
            for idx in indices:
                lat_idx, lon_idx = np.unravel_index(idx, lat_fy.shape)
                latitudes.append(lat_fy[lat_idx, lon_idx])
                longitudes.append(lon_fy[lat_idx, lon_idx])
                # cth_value = cth_data[lat_idx, lon_idx]*0.001
                # 将每个匹配点的CTH值添加到列表中
                cth_values.append(cth_data[lat_idx, lon_idx] * 0.001)
                cbh_values.append(cbh_data[lat_idx, lon_idx] * 0.001)
                    # print(
                #     f"Found point at lat_fy[{lat_idx}][{lon_idx}] = {lat_fy[lat_idx, lon_idx]}, lon_fy[{lat_idx}][{lon_idx}] = {lon_fy[lat_idx, lon_idx]}")
                # print(f"CTH value at lat_fy[{lat_idx}][{lon_idx}] = {cth_value}")
        else:
            # 如果没有找到匹配，添加NaN作为占位符
            latitudes.append(np.nan)
            longitudes.append(np.nan)
            cth_values.append(np.nan)
            cbh_values.append(np.nan)

    fy_data = {
        'Latitude': latitudes,
        'Longitude': longitudes,
        'CTH': cth_values,
        'CBH': cbh_values
    }

    df = pd.DataFrame(fy_data)

    # 过滤掉包含NaN值的行
    # df.dropna(inplace=True)
    # 按照纬度排序
    df_sorted = df.sort_values(by='Latitude')

    # 获取排序后的纬度和CTH
    latitude_sorted = df_sorted['Latitude']
    cth_sorted = df_sorted['CTH']
    cbh_sorted = df_sorted['CBH']
    # # 创建折线图
    # plt.figure(figsize=(10, 6))
    # plt.plot(latitude_sorted, cth_sorted, color='blue', label='CTH vs Latitude')
    # plt.xlabel('Latitude')
    # plt.ylabel('Cloud Top Height (m)')
    # plt.title('Latitude vs. Mean Cloud Top Height')
    # plt.savefig('fy201902220700', bbox_inches='tight')
    # plt.grid(True)
    # plt.legend()  # 显示图例
    # # plt.show()

    # 绘制剖面图
    fig = plt.figure(figsize=(9, 5))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_title(f'From {start_datetime} to {end_datetime}', fontsize='small')

    # 绘制剖面图
    shade_plt = ax.pcolormesh(lat_t, hgt, rad_refl_t, vmin=-20, vmax=30)
    line_plt = ax.plot(latitude_sorted, cth_sorted, color='blue', label='CTH vs Latitude')
    line2_plt= ax.plot(latitude_sorted, cbh_sorted, color='red', label='CBH vs Latitude')
    ax.set_ylim(0, 20)
    ax.set_xlim(-11,-6)
    ax.set_xlabel('Latitude')
    ax.set_ylabel('Height (km)')
    cbar = fig.colorbar(shade_plt, ax=ax, shrink=0.8, extend="both")
    fig.savefig('cloudsatmatchfy201902220700', bbox_inches='tight')
    plt.show()
    plt.close(fig)










